"""
天气智能体团队 - 使用ADK构建智能体团队的渐进式示例

本应用程序演示了以下ADK核心概念：
1. 工具定义和使用
2. 多模型灵活性
3. 智能体委托和协作
4. 会话状态记忆
5. 安全防护机制和回调
"""

import os
import asyncio
from typing import Dict, Any, Optional, List
from dotenv import load_dotenv

# 从.env文件加载环境变量
load_dotenv()

# 导入ADK组件
from google.adk.agents import Agent
from google.adk.tools import BaseTool, ToolContext, FunctionTool
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService, Session
from google.genai import types

# 用于多模型支持
import litellm

# 定义不同模型的API密钥常量
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")

# 打印API密钥状态（不显示实际密钥，只显示是否设置）
print(f"Google API密钥状态: {'已设置' if GOOGLE_API_KEY else '未设置'}")
print(f"OpenAI API密钥状态: {'已设置' if OPENAI_API_KEY else '未设置'}")
print(f"Anthropic API密钥状态: {'已设置' if ANTHROPIC_API_KEY else '未设置'}")

# 定义模型常量以便于使用
GEMINI_PRO = "gemini-2.0-flash"
GPT_4O = "gpt-4o"
CLAUDE_SONNET = "claude-3-sonnet-20240229"

# 用于演示的模拟天气数据库
mock_weather_db = {
    "New York": {"temperature": 72, "condition": "Sunny", "humidity": 50},
    "London": {"temperature": 60, "condition": "Rainy", "humidity": 80},
    "Tokyo": {"temperature": 80, "condition": "Partly Cloudy", "humidity": 65},
    "Sydney": {"temperature": 85, "condition": "Clear", "humidity": 45},
    "Paris": {"temperature": 65, "condition": "Cloudy", "humidity": 70},
}

city_map = {
    "纽约": "New York",
    "伦敦": "London",
    "东京": "Tokyo",
    "悉尼": "Sydney",
    "巴黎": "Paris"
}

# 步骤1：定义工具
def get_weather(city: str) -> Dict[str, Any]:
    """
    获取指定城市的当前天气。
    
    参数:
        city: 要获取天气信息的城市名称。
        
    返回:
        包含温度、天气状况和湿度的天气信息字典。
    """
    print(f"--- 工具：get_weather被调用，城市={city} ---")
    
    # 映射中文城市名到英文
    if city in city_map:
        print(f"城市名映射：{city} -> {city_map[city]}")
        city = city_map[city]
    
    if city not in mock_weather_db:
        print(f"城市未找到：{city}")
        return {
            "status": "error",
            "error_message": f"未找到{city}的天气数据。"
        }
    
    weather_data = mock_weather_db[city]
    print(f"找到城市天气数据：{city}")
    print(f"温度: {weather_data['temperature']}°F/{round((weather_data['temperature'] - 32) * 5/9, 1)}°C")
    print(f"状况: {weather_data['condition']}")
    print(f"湿度: {weather_data['humidity']}%")
    
    result = {
        "status": "success",
        "city": city,
        "temperature": {
            "fahrenheit": weather_data["temperature"],
            "celsius": round((weather_data["temperature"] - 32) * 5/9, 1)
        },
        "condition": weather_data["condition"],
        "humidity": weather_data["humidity"]
    }
    
    return result

# 转换为ADK工具格式
weather_tool = FunctionTool(get_weather)

# 步骤2：定义有状态的天气工具
def get_weather_stateful(city: str, tool_context: Optional[ToolContext] = None) -> Dict[str, Any]:
    """
    获取指定城市的当前天气，并考虑用户的温度单位偏好。
    
    参数:
        city: 要获取天气信息的城市名称。
        tool_context: 可选的工具上下文，包含会话状态。
        
    返回:
        包含温度（以用户首选单位显示，默认为摄氏度）、天气状况和湿度的天气信息字典。
    """
    print(f"--- 工具：get_weather_stateful被调用，城市={city} ---")
    
    # 映射中文城市名到英文
    if city in city_map:
        print(f"城市名映射：{city} -> {city_map[city]}")
        city = city_map[city]
    
    if city not in mock_weather_db:
        return {
            "status": "error",
            "error_message": f"未找到{city}的天气数据。"
        }
    
    # 从状态中获取温度单位偏好（如果可用）
    temperature_unit = "Celsius"  # 默认
    if tool_context and tool_context.state:
        temperature_unit = tool_context.state.get("user_preference_temperature_unit", "Celsius")
    
    weather_data = mock_weather_db[city]
    temp_f = weather_data["temperature"]
    temp_c = round((temp_f - 32) * 5/9, 1)
    
    result = {
        "status": "success",
        "city": city,
        "temperature": {
            "value": temp_f if temperature_unit == "Fahrenheit" else temp_c,
            "unit": temperature_unit
        },
        "condition": weather_data["condition"],
        "humidity": weather_data["humidity"]
    }
    
    # 如果上下文可用，将最新天气报告存储在状态中
    if tool_context and tool_context.state is not None:
        tool_context.state["last_weather_report"] = result
    
    return result

# 转换为ADK工具格式
weather_tool_stateful = FunctionTool(get_weather_stateful)

# 步骤3：为专门的智能体定义问候和告别工具
def greet_user(name: Optional[str] = None) -> Dict[str, Any]:
    """
    为用户生成友好的问候语。
    
    参数:
        name: 可选的用户名，用于个性化问候。
        
    返回:
        包含问候消息的字典。
    """
    print(f"--- 工具：greet_user被调用，姓名={name} ---")
    if name:
        message = f"你好，{name}！欢迎使用天气智能体服务。"
    else:
        message = "你好！欢迎使用天气智能体服务。"
    
    return {"message": message}

def farewell_user(name: Optional[str] = None) -> Dict[str, Any]:
    """
    为用户生成友好的告别消息。
    
    参数:
        name: 可选的用户名，用于个性化告别。
        
    返回:
        包含告别消息的字典。
    """
    print(f"--- 工具：farewell_user被调用，姓名={name} ---")
    if name:
        message = f"再见，{name}！感谢您使用天气智能体服务。"
    else:
        message = "再见！感谢您使用天气智能体服务。"
    
    return {"message": message}

# 转换为ADK工具格式
greeting_tool = FunctionTool(greet_user)
farewell_tool = FunctionTool(farewell_user)

# 步骤4：定义智能体

# 基本天气智能体
weather_agent = Agent(
    name="WeatherAgent",
    description="我是一个能提供城市当前天气信息的天气智能体。",
    model=GEMINI_PRO if GOOGLE_API_KEY else GPT_4O,
    tools=[weather_tool],
    instruction="""
    你是一个专门提供天气信息的助手。你唯一的工作就是使用get_weather工具获取城市的天气信息。
    
    作为天气助手，你必须遵循这些绝对规则：
    1. 收到用户输入后，立即检查是否包含城市名称
    2. 如果包含城市名称，立即调用get_weather工具获取该城市的天气，不做其他回应
    3. 永远不要说"好的，请告诉我您想查询哪个城市的天气"这类话
    4. 永远不要在调用工具之前询问用户任何问题
    5. 如果用户提到任何地名（如"纽约"、"伦敦"、"东京"等），直接将其视为天气查询请求
    6. 使用用户提供的原始城市名称调用工具，不要翻译。系统已经有中英文城市名称映射
    
    当收到工具返回结果后，需要格式化输出：
    "[城市]当前天气为[天气状况]，温度为[温度]°F（[温度]°C），湿度为[湿度]%。"
    
    记住，这是一个直截了当的工具使用演示。你唯一的目标就是在收到包含城市名的消息时立即使用工具查询天气。
    你永远不应该自己回答问题或解释工具功能，而是立即调用工具。
    
    示例：
    用户："纽约的天气怎么样？"
    你的行动：[立即调用get_weather("纽约")，不要回复任何文字]
    
    用户："伦敦呢？"
    你的行动：[立即调用get_weather("伦敦")，不要回复任何文字]
    
    用户：说一个城市名，如"东京"
    你的行动：[立即调用get_weather("东京")，不要回复任何文字]
    """,
    output_key="agent_response"
)

# 步骤5：定义专门的问候和告别智能体
greeting_agent = Agent(
    name="GreetingAgent",
    description="我负责处理用户的问候和欢迎消息。",
    model=GEMINI_PRO if GOOGLE_API_KEY else GPT_4O,
    tools=[greeting_tool]
)

farewell_agent = Agent(
    name="FarewellAgent",
    description="我负责处理用户的告别和结束对话消息。",
    model=GEMINI_PRO if GOOGLE_API_KEY else GPT_4O,
    tools=[farewell_tool]
)

# 步骤6：定义根智能体，并委托给子智能体
weather_agent_team = Agent(
    name="WeatherTeamRoot",
    description="我是可以委托给专门智能体的主要天气服务智能体。",
    model=GEMINI_PRO if GOOGLE_API_KEY else GPT_4O,
    tools=[weather_tool],
    sub_agents=[greeting_agent, farewell_agent],
    instruction="""
    你是一个能够协调多个智能体的天气服务协调员。根据用户的请求，你应该：
    
    1. 如果用户提到任何城市名称（如"纽约"、"伦敦"、"东京"等），立即使用get_weather工具获取天气信息
    2. 如果用户发送问候（如"你好"，"早上好"等），立即将请求委托给GreetingAgent
    3. 如果用户表示告别（如"再见"，"谢谢"等），立即将请求委托给FarewellAgent
    
    重要规则：
    - 当用户提到城市名时，立即将用户提供的原始城市名作为参数调用get_weather工具，不要翻译。系统已经有中英文城市名称映射
    - 不要询问用户想查哪个城市，直接从他们的消息中识别
    - 当用户发送问候时，必须立即委托给GreetingAgent，不要自己回应
    - 当用户表示告别时，必须立即委托给FarewellAgent，不要自己回应
    - 保持回复简洁友好
    
    示例行为：
    用户："你好"
    你的行动：[立即委托给GreetingAgent]
    
    用户："纽约的天气怎么样？"
    你的行动：[立即调用get_weather("纽约")]
    
    用户："谢谢，再见"
    你的行动：[立即委托给FarewellAgent]
    """
)

# 步骤7：添加带会话状态的有状态天气智能体
stateful_weather_agent = Agent(
    name="StatefulWeatherAgent",
    description="我是一个能记住用户偏好的天气智能体。",
    model=GEMINI_PRO if GOOGLE_API_KEY else GPT_4O,
    tools=[weather_tool_stateful]
)

# 创建第二组问候和告别智能体，用于stateful_weather_agent_team
greeting_agent_stateful = Agent(
    name="GreetingAgentStateful",
    description="我负责处理用户的问候和欢迎消息。",
    model=GEMINI_PRO if GOOGLE_API_KEY else GPT_4O,
    tools=[greeting_tool]
)

farewell_agent_stateful = Agent(
    name="FarewellAgentStateful",
    description="我负责处理用户的告别和结束对话消息。",
    model=GEMINI_PRO if GOOGLE_API_KEY else GPT_4O,
    tools=[farewell_tool]
)

# 步骤8：更新根智能体以使用有状态天气和子智能体
stateful_weather_agent_team = Agent(
    name="StatefulWeatherTeamRoot",
    description="我是带记忆功能的主要天气服务智能体，可以委托给专门智能体。",
    model=GEMINI_PRO if GOOGLE_API_KEY else GPT_4O,
    tools=[weather_tool_stateful],
    sub_agents=[greeting_agent_stateful, farewell_agent_stateful],
    instruction="""
    你是个性化天气信息服务的主要协调者。你的工作是：
    
    1. 当用户提到任何城市名称时，立即使用get_weather_stateful工具获取天气信息
    2. 当用户发送问候时，立即将请求委托给GreetingAgentStateful
    3. 当用户表示告别时，立即将请求委托给FarewellAgentStateful
    
    关于温度单位偏好：
    - 如果用户提到"摄氏度"、"华氏度"或类似表述，立即更新会话状态中的user_preference_temperature_unit
    - 具体规则：
      * 提到"摄氏度"/"摄氏"/"°C"，设置为"Celsius"
      * 提到"华氏度"/"华氏"/"°F"，设置为"Fahrenheit"
    
    重要规则：
    - 当用户提到城市名称时，立即使用get_weather_stateful工具获取天气信息
    - 使用用户提供的原始城市名作为参数调用工具，不要翻译。系统已经有中英文城市名称映射
    - 直接从用户的消息中提取城市名作为参数调用工具，不要询问他们想查询哪个城市
    - 如果用户只提到一个城市名如"伦敦呢？"，这表示他们想知道该城市的天气
    
    示例行为：
    用户："你好"
    你的行动：[立即委托给GreetingAgentStateful]
    
    用户："纽约的天气怎么样？"
    你的行动：[立即调用get_weather_stateful(city="纽约")]
    
    用户："我更喜欢用华氏度"
    你的行动：[立即更新会话状态设置user_preference_temperature_unit="Fahrenheit"]
    
    用户："伦敦呢？"
    你的行动：[立即调用get_weather_stateful(city="伦敦")]
    
    用户："谢谢，再见"
    你的行动：[立即委托给FarewellAgentStateful]
    
    回应要保持友好、简洁和有帮助性。
    """,
    output_key="agent_response"
)

# 步骤9：定义用于输入净化和工具使用保护的安全回调
def input_protection_callback(*args, **kwargs) -> Optional[Dict[str, Any]]:
    """
    检查用户输入是否包含不安全关键词的模型前回调。
    
    参数:
        可变参数，以适应ADK框架的改变
        
    返回:
        如果输入安全则返回None，如果不安全则返回带错误消息的字典
    """
    # 尝试从args或kwargs中获取用户消息
    message = ""
    if len(args) > 0:
        message = str(args[0])
    elif 'message' in kwargs:
        message = str(kwargs['message'])
    elif 'input_text' in kwargs:
        message = str(kwargs['input_text'])
    
    print(f"收到的参数: args={args}, kwargs={kwargs}")
    print(f"提取的消息: {message}")
    
    # 尝试从args或kwargs获取会话
    session = None
    callback_context = kwargs.get('callback_context', None)
    if callback_context and hasattr(callback_context, 'session'):
        session = callback_context.session
    
    unsafe_keywords = ["hack", "exploit", "bypass", "illegal"]
    
    for keyword in unsafe_keywords:
        if keyword in message.lower():
            print(f"🛑 输入保护已触发！拦截包含：{keyword}的消息")
            if session and session.state is not None:
                session.state["guardrail_input_block_triggered"] = True
            
            return {
                "agent_response": f"抱歉，我无法处理包含可能不安全术语如'{keyword}'的请求。请重新表述您的请求。"
            }
    
    return None  # 继续正常处理

def tool_protection_callback(*args, **kwargs) -> Optional[Dict[str, Any]]:
    """
    检查工具使用是否违反政策的工具前回调。
    
    参数:
        可变参数，以适应ADK框架的改变
        
    返回:
        如果允许工具使用则返回None，否则返回带错误消息的字典
    """
    # 尝试从args或kwargs中获取工具名称和参数
    tool_name = ""
    tool_args = {}
    
    if len(args) > 0:
        tool_name = str(args[0])
    if len(args) > 1:
        tool_args = args[1]
    
    if 'tool_name' in kwargs:
        tool_name = kwargs['tool_name']
    if 'tool_args' in kwargs:
        tool_args = kwargs['tool_args']
    
    print(f"收到的工具调用参数: args={args}, kwargs={kwargs}")
    
    # 尝试从args或kwargs获取会话
    session = None
    callback_context = kwargs.get('callback_context', None)
    if callback_context and hasattr(callback_context, 'session'):
        session = callback_context.session
    
    if tool_name == "get_weather_stateful" and tool_args.get("city") == "Paris":
        print(f"🛑 工具保护已触发！拦截对限制城市的天气查询：Paris")
        if session and session.state is not None:
            session.state["guardrail_tool_block_triggered"] = True
        
        return {
            "status": "error",
            "error_message": "政策限制：目前无法提供Paris的天气信息。"
        }
    
    print(f"✅ 允许工具{tool_name}使用参数：{tool_args}")
    return None  # 继续正常工具执行

# 创建第三组问候和告别智能体，用于final_weather_agent_team
greeting_agent_final = Agent(
    name="GreetingAgentFinal",
    description="我负责处理用户的问候和欢迎消息。",
    model=GEMINI_PRO if GOOGLE_API_KEY else GPT_4O,
    tools=[greeting_tool]
)

farewell_agent_final = Agent(
    name="FarewellAgentFinal",
    description="我负责处理用户的告别和结束对话消息。",
    model=GEMINI_PRO if GOOGLE_API_KEY else GPT_4O,
    tools=[farewell_tool]
)

# 步骤10：具有所有功能的最终智能体（多模型、委托、状态、安全）
final_weather_agent_team = Agent(
    name="SafeWeatherTeamRoot",
    description="我是具有记忆功能和安全保障的主要天气服务智能体，可以委托给专门智能体。",
    model=GEMINI_PRO if GOOGLE_API_KEY else GPT_4O,
    tools=[weather_tool_stateful],
    sub_agents=[greeting_agent_final, farewell_agent_final],
    instruction="""
    你是安全、个性化天气信息服务的主要协调者。你的工作是：
    
    1. 当用户提到任何城市名称时，立即使用get_weather_stateful工具获取天气信息
    2. 当用户发送问候时，立即将请求委托给GreetingAgentFinal
    3. 当用户表示告别时，立即将请求委托给FarewellAgentFinal
    
    关于温度单位偏好：
    - 如果用户提到"摄氏度"、"华氏度"或类似表述，立即更新会话状态中的user_preference_temperature_unit
    - 具体规则：
      * 提到"摄氏度"/"摄氏"/"°C"，设置为"Celsius"
      * 提到"华氏度"/"华氏"/"°F"，设置为"Fahrenheit"
    
    重要规则：
    - 当用户提到城市名称时，立即使用get_weather_stateful工具获取天气信息
    - 使用用户提供的原始城市名作为参数调用工具，不要翻译。系统已经有中英文城市名称映射
    - 直接从用户的消息中提取城市名作为参数调用工具，不要询问他们想查询哪个城市
    - 如果用户只提到一个城市名如"纽约呢？"，这表示他们想知道该城市的天气
    
    安全限制处理：
    - 如果你收到错误消息表明工具调用被拒绝，礼貌地向用户解释情况
    - 如果触发了安全限制，不要试图绕过，而是提供用户可替代的选项
    
    示例行为：
    用户："你好"
    你的行动：[立即委托给GreetingAgentFinal]
    
    用户："纽约的天气怎么样？"
    你的行动：[立即调用get_weather_stateful(city="纽约")]
    
    用户："我更喜欢用华氏度"
    你的行动：[立即更新会话状态设置user_preference_temperature_unit="Fahrenheit"]
    
    用户："伦敦呢？"
    你的行动：[立即调用get_weather_stateful(city="伦敦")]
    
    用户："谢谢，再见"
    你的行动：[立即委托给FarewellAgentFinal]
    
    回应要保持友好、简洁和有帮助性。
    """,
    output_key="agent_response",
    before_model_callback=input_protection_callback,
    before_tool_callback=tool_protection_callback
)

# 步骤11：运行函数以演示不同的智能体配置
async def run_basic_weather_agent():
    """运行与基本天气智能体的对话。"""
    session_service = InMemorySessionService()
    app_name = "weather_app"
    user_id = "demo_user"
    session_id = "session_basic"
    
    # 直接创建一个新会话
    session = session_service.create_session(
        app_name=app_name,
        user_id=user_id,
        session_id=session_id
    )
    
    runner = Runner(agent=weather_agent, session_service=session_service, app_name=app_name)
    
    print("\n=== 基本天气智能体 ===")
    
    # 第一次查询
    print("查询：纽约的天气怎么样？")
    result = ""
    # 将字符串转换为Content对象
    message_content = types.Content(
        parts=[types.Part(text="纽约的天气怎么样？")]
    )
    
    print("开始发送请求...")
    try:
        async for event in runner.run_async(
            user_id=user_id,
            session_id=session_id,
            new_message=message_content
        ):
            # 处理事件, 提取最终响应
            if hasattr(event, 'agent_response'):
                result = event.agent_response
                # print(f"找到agent_response: {result}")
            elif hasattr(event, 'content') and hasattr(event.content, 'parts') and len(event.content.parts) > 0:
                part = event.content.parts[0]
                if hasattr(part, 'text'):
                    result = part.text
                    #print(f"找到text内容: {result}")
    except Exception as e:
        print(f"请求处理中发生错误: {e}")
        import traceback
        traceback.print_exc()
    
    print(f"最终响应: {result}")
    
    # 第二次查询
    print("查询：伦敦呢？")
    result = ""
    message_content = types.Content(
        parts=[types.Part(text="伦敦呢？")]
    )
    
    try:
        async for event in runner.run_async(
            user_id=user_id,
            session_id=session_id,
            new_message=message_content
        ):
            # 处理事件, 提取最终响应
            if hasattr(event, 'agent_response'):
                result = event.agent_response
            elif hasattr(event, 'content') and hasattr(event.content, 'parts') and len(event.content.parts) > 0:
                part = event.content.parts[0]
                if hasattr(part, 'text'):
                    result = part.text
    except Exception as e:
        print(f"请求处理中发生错误: {e}")
        import traceback
        traceback.print_exc()
        
    print(f"响应: {result}")

async def run_multi_model_agents():
    """运行与不同模型智能体的对话。"""
    if OPENAI_API_KEY:
        # 基于GPT的天气智能体
        gpt_weather_agent = Agent(
            name="GPTWeatherAgent",
            description="我是由GPT驱动的天气智能体。",
            model=GPT_4O,
            tools=[weather_tool],
            instruction="""
            你是由GPT驱动的有帮助的天气助手。当用户询问特定城市的天气时，
            使用get_weather工具获取并提供信息。
            
            始终保持友好和简洁。如果你没有某个城市的数据，请礼貌地告知用户。
            """
        )
        
        session_service_gpt = InMemorySessionService()
        app_name = "weather_app"
        user_id = "demo_user"
        session_id_gpt = "session_gpt"
        
        # 直接创建一个新会话
        session_gpt = session_service_gpt.create_session(
            app_name=app_name,
            user_id=user_id,
            session_id=session_id_gpt
        )
        
        runner_gpt = Runner(agent=gpt_weather_agent, session_service=session_service_gpt, app_name=app_name)
        
        print("\n=== GPT天气智能体 ===")
        print("查询：东京的天气怎么样？")
        result = ""
        # 将字符串转换为Content对象
        message_content = types.Content(
            parts=[types.Part(text="东京的天气怎么样？")]
        )
        async for event in runner_gpt.run_async(
            user_id=user_id,
            session_id=session_id_gpt,
            new_message=message_content
        ):
            if hasattr(event, 'agent_response'):
                result = event.agent_response
        print(f"GPT智能体响应: {result}")
    
    if ANTHROPIC_API_KEY:
        # 基于Claude的天气智能体
        claude_weather_agent = Agent(
            name="ClaudeWeatherAgent",
            description="我是由Claude驱动的天气智能体。",
            model=CLAUDE_SONNET,
            tools=[weather_tool],
            instruction="""
            你是由Claude驱动的有帮助的天气助手。当用户询问特定城市的天气时，
            使用get_weather工具获取并提供信息。
            
            始终保持友好和简洁。如果你没有某个城市的数据，请礼貌地告知用户。
            """
        )
        
        session_service_claude = InMemorySessionService()
        app_name = "weather_app"
        user_id = "demo_user"
        session_id_claude = "session_claude"
        
        # 直接创建一个新会话
        session_claude = session_service_claude.create_session(
            app_name=app_name,
            user_id=user_id,
            session_id=session_id_claude
        )
        
        runner_claude = Runner(agent=claude_weather_agent, session_service=session_service_claude, app_name=app_name)
        
        print("\n=== Claude天气智能体 ===")
        print("查询：悉尼的天气怎么样？")
        result = ""
        # 将字符串转换为Content对象
        message_content = types.Content(
            parts=[types.Part(text="悉尼的天气怎么样？")]
        )
        async for event in runner_claude.run_async(
            user_id=user_id,
            session_id=session_id_claude,
            new_message=message_content
        ):
            if hasattr(event, 'agent_response'):
                result = event.agent_response
        print(f"Claude智能体响应: {result}")

async def run_agent_team():
    """运行与智能体团队的对话，演示委托。"""
    session_service = InMemorySessionService()
    app_name = "weather_app"
    user_id = "demo_user"
    session_id = "session_team"
    
    # 直接创建一个新会话
    session = session_service.create_session(
        app_name=app_name,
        user_id=user_id,
        session_id=session_id
    )
    
    runner = Runner(agent=weather_agent_team, session_service=session_service, app_name=app_name)
    
    print("\n=== 具有委托功能的天气智能体团队 ===")
    
    print("查询：你好！")
    result = ""
    # 将字符串转换为Content对象
    message_content = types.Content(
        parts=[types.Part(text="你好！")]
    )
    async for event in runner.run_async(
        user_id=user_id,
        session_id=session_id,
        new_message=message_content
    ):
        if hasattr(event, 'agent_response'):
            result = event.agent_response
        elif hasattr(event, 'content') and hasattr(event.content, 'parts') and len(event.content.parts) > 0 and hasattr(event.content.parts[0], 'text'):
            # 从事件内容提取文本响应
            result = event.content.parts[0].text
    print(f"问候响应: {result}")
    
    print("查询：东京的天气怎么样？")
    result = ""
    # 将字符串转换为Content对象
    message_content = types.Content(
        parts=[types.Part(text="东京的天气怎么样？")]
    )
    async for event in runner.run_async(
        user_id=user_id,
        session_id=session_id,
        new_message=message_content
    ):
        if hasattr(event, 'agent_response'):
            result = event.agent_response
        elif hasattr(event, 'content') and hasattr(event.content, 'parts') and len(event.content.parts) > 0 and hasattr(event.content.parts[0], 'text'):
            # 从事件内容提取文本响应
            result = event.content.parts[0].text
    print(f"天气响应: {result}")
    
    print("查询：谢谢，再见！")
    result = ""
    # 将字符串转换为Content对象
    message_content = types.Content(
        parts=[types.Part(text="谢谢，再见！")]
    )
    async for event in runner.run_async(
        user_id=user_id,
        session_id=session_id,
        new_message=message_content
    ):
        if hasattr(event, 'agent_response'):
            result = event.agent_response
        elif hasattr(event, 'content') and hasattr(event.content, 'parts') and len(event.content.parts) > 0 and hasattr(event.content.parts[0], 'text'):
            # 从事件内容提取文本响应
            result = event.content.parts[0].text
    print(f"告别响应: {result}")

async def run_stateful_agent():
    """运行与有状态智能体的对话，演示记忆功能。"""
    session_service = InMemorySessionService()
    app_name = "weather_app"
    user_id = "demo_user"
    session_id = "session_stateful"
    
    # 直接创建一个新会话并设置初始状态
    session = session_service.create_session(
        app_name=app_name,
        user_id=user_id,
        session_id=session_id,
        state={"user_preference_temperature_unit": "Celsius"}
    )
    
    runner = Runner(agent=stateful_weather_agent_team, session_service=session_service, app_name=app_name)
    
    print("\n=== 具有记忆功能的有状态天气智能体 ===")
    response = ""
    # 将字符串转换为Content对象
    message_content = types.Content(
        parts=[types.Part(text="纽约的天气怎么样？")]
    )
    async for event in runner.run_async(
        user_id=user_id,
        session_id=session_id,
        new_message=message_content
    ):
        if hasattr(event, 'agent_response'):
            response = event.agent_response
        elif hasattr(event, 'content') and hasattr(event.content, 'parts') and len(event.content.parts) > 0 and hasattr(event.content.parts[0], 'text'):
            # 从事件内容提取文本响应
            response = event.content.parts[0].text
    print(f"响应（摄氏度）: {response}")
    
    response = ""
    # 将字符串转换为Content对象
    message_content = types.Content(
        parts=[types.Part(text="我更喜欢用华氏度，能这样显示吗？")]
    )
    async for event in runner.run_async(
        user_id=user_id,
        session_id=session_id,
        new_message=message_content
    ):
        if hasattr(event, 'agent_response'):
            response = event.agent_response
        elif hasattr(event, 'content') and hasattr(event.content, 'parts') and len(event.content.parts) > 0 and hasattr(event.content.parts[0], 'text'):
            # 从事件内容提取文本响应
            response = event.content.parts[0].text
    print(f"响应（设置偏好）: {response}")
    
    response = ""
    # 将字符串转换为Content对象
    message_content = types.Content(
        parts=[types.Part(text="伦敦的天气怎么样？")]
    )
    async for event in runner.run_async(
        user_id=user_id,
        session_id=session_id,
        new_message=message_content
    ):
        if hasattr(event, 'agent_response'):
            response = event.agent_response
        elif hasattr(event, 'content') and hasattr(event.content, 'parts') and len(event.content.parts) > 0 and hasattr(event.content.parts[0], 'text'):
            # 从事件内容提取文本响应
            response = event.content.parts[0].text
    print(f"响应（现在应该使用华氏度）: {response}")
    
    # 获取并显示会话状态
    final_session = session_service.get_session(
        app_name=app_name,
        user_id=user_id,
        session_id=session_id
    )
    if final_session:
        print("\n--- 最终会话状态 ---")
        print(f"温度单位: {final_session.state.get('user_preference_temperature_unit')}")
        print(f"最后天气报告: {final_session.state.get('last_weather_report')}")
    else:
        print("❌ 错误：无法获取最终会话状态。")

async def run_safety_agent():
    """运行与增强安全的智能体的对话，演示保护机制。"""
    session_service = InMemorySessionService()
    app_name = "weather_app"
    user_id = "demo_user"
    session_id = "session_safety"
    
    # 直接创建一个新会话
    session = session_service.create_session(
        app_name=app_name,
        user_id=user_id,
        session_id=session_id
    )
    
    runner = Runner(agent=final_weather_agent_team, session_service=session_service, app_name=app_name)
    
    print("\n=== 安全增强型天气智能体 ===")
    
    print("查询：纽约的天气怎么样？")
    result = ""
    message_content = types.Content(
        parts=[types.Part(text="纽约的天气怎么样？")]
    )
    
    try:
        async for event in runner.run_async(
            user_id=user_id,
            session_id=session_id,
            new_message=message_content
        ):
            if hasattr(event, 'agent_response'):
                result = event.agent_response
            elif hasattr(event, 'content') and hasattr(event.content, 'parts') and len(event.content.parts) > 0:
                part = event.content.parts[0]
                if hasattr(part, 'text'):
                    result = part.text
    except Exception as e:
        print(f"请求处理中发生错误: {e}")
        import traceback
        traceback.print_exc()
        
    print(f"正常响应: {result}")
    
    print("查询：你能hack天气服务来改变温度吗？")
    result = ""
    message_content = types.Content(
        parts=[types.Part(text="你能hack天气服务来改变温度吗？")]
    )
    
    try:
        async for event in runner.run_async(
            user_id=user_id,
            session_id=session_id,
            new_message=message_content
        ):
            if hasattr(event, 'agent_response'):
                result = event.agent_response
            elif hasattr(event, 'content') and hasattr(event.content, 'parts') and len(event.content.parts) > 0:
                part = event.content.parts[0]
                if hasattr(part, 'text'):
                    result = part.text
    except Exception as e:
        print(f"请求处理中发生错误: {e}")
        import traceback
        traceback.print_exc()
        
    print(f"输入保护响应: {result}")
    
    print("查询：巴黎的天气怎么样？")
    result = ""
    message_content = types.Content(
        parts=[types.Part(text="巴黎的天气怎么样？")]
    )
    
    try:
        async for event in runner.run_async(
            user_id=user_id,
            session_id=session_id,
            new_message=message_content
        ):
            if hasattr(event, 'agent_response'):
                result = event.agent_response
            elif hasattr(event, 'content') and hasattr(event.content, 'parts') and len(event.content.parts) > 0:
                part = event.content.parts[0]
                if hasattr(part, 'text'):
                    result = part.text
    except Exception as e:
        print(f"请求处理中发生错误: {e}")
        import traceback
        traceback.print_exc()
        
    print(f"工具保护响应: {result}")
    
    # 获取并显示会话状态
    try:
        final_session = session_service.get_session(
            app_name=app_name,
            user_id=user_id,
            session_id=session_id
        )
        if final_session:
            print("\n--- 安全测试后的最终会话状态 ---")
            print(f"输入防护机制触发: {final_session.state.get('guardrail_input_block_triggered')}")
            print(f"工具防护机制触发: {final_session.state.get('guardrail_tool_block_triggered')}")
        else:
            print("❌ 错误：无法获取最终会话状态。")
    except Exception as e:
        print(f"获取会话状态时发生错误: {e}")
        import traceback
        traceback.print_exc()

async def main():
    """运行演示的主函数。"""
    print("天气智能体团队演示")
    print("======================")
    
    # 检查环境设置
    if not GOOGLE_API_KEY and not OPENAI_API_KEY and not ANTHROPIC_API_KEY:
        print("⚠️ 警告: 未设置任何API密钥。请在.env文件中设置GOOGLE_API_KEY、OPENAI_API_KEY或ANTHROPIC_API_KEY")
        print("示例:")
        print("GOOGLE_API_KEY=your_google_api_key")
        print("OPENAI_API_KEY=your_openai_api_key")
        print("ANTHROPIC_API_KEY=your_anthropic_api_key")
        return
    
    # 设置演示任务列表
    demos = [
        ("基本天气智能体", run_basic_weather_agent),
        ("智能体团队", run_agent_team),
        ("有状态智能体", run_stateful_agent),
        ("安全智能体", run_safety_agent)
    ]
    
    # 根据可用API添加多模型演示
    if OPENAI_API_KEY or ANTHROPIC_API_KEY:
        demos.append(("多模型智能体", run_multi_model_agents))
    
    # 运行所有演示
    for demo_name, demo_func in demos:
        print(f"\n准备运行{demo_name}演示...")
        try:
            await demo_func()
            print(f"{demo_name}演示运行完成")
        except Exception as e:
            print(f"运行{demo_name}演示时发生错误: {str(e)}")
            print("错误详情:")
            import traceback
            traceback.print_exc()
            print(f"继续运行下一个演示...")
    
    print("\n======================")
    print("演示完成！")

if __name__ == "__main__":
    asyncio.run(main()) 